#include "task.h"
#include "uart.h"

static struct task tasks[MAX_TASKS];  // 任务数组
static int current_task = 0;          // 当前任务
static int num_tasks = 0;             // 任务数量

// 保存上下文并切换到新任务
static void switch_to(int next_task);

// 汇编函数声明
extern void task_switch(uint64_t* prev_sp, uint64_t* next_sp);

void task_init(void) {
    // 初始化任务数组
    for(int i = 0; i < MAX_TASKS; i++) {
        tasks[i].status = TASK_FINISHED;
    }
    uart_puts("Task system initialized\n");
}

int task_create(const char* name, task_func_t func) {
    if(num_tasks >= MAX_TASKS) {
        uart_puts("Cannot create more tasks!\n");
        return -1;
    }

    int id = num_tasks++;
    struct task* t = &tasks[id];
    
    // 设置任务信息
    for(int i = 0; i < 15 && name[i]; i++) {
        t->name[i] = name[i];
    }
    t->name[15] = '\0';
    
    // 初始化任务栈
    t->sp = (uint64_t)&t->stack[STACK_SIZE - 16];
    t->status = TASK_READY;
    
    // 设置任务初始状态
    uint64_t* sp = (uint64_t*)t->sp;
    sp[0] = (uint64_t)func;     // ra 返回地址设为任务函数
    sp[-1] = 0;                 // 之前的帧指针
    t->sp = (uint64_t)&sp[-2];  // 更新栈指针
    
    uart_puts("Created task: ");
    uart_puts(t->name);
    uart_puts("\n");
    
    return id;
}

void task_yield(void) {
    // 找到下一个就绪任务
    int next = (current_task + 1) % num_tasks;
    while(next != current_task) {
        if(tasks[next].status == TASK_READY) {
            break;
        }
        next = (next + 1) % num_tasks;
    }
    
    if(next != current_task && tasks[next].status == TASK_READY) {
        int prev = current_task;
        current_task = next;
        tasks[next].status = TASK_RUNNING;
        tasks[prev].status = TASK_READY;
        switch_to(next);
    }
}

static void switch_to(int next_task) {
    task_switch(&tasks[current_task].sp, &tasks[next_task].sp);
}

void task_start(void) {
    if(num_tasks > 0) {
        current_task = 0;
        tasks[0].status = TASK_RUNNING;
        uint64_t sp = tasks[0].sp;
        
        // 加载第一个任务的上下文并开始执行
        asm volatile(
            "mv sp, %0\n"
            "ld ra, 0(sp)\n"
            "addi sp, sp, 8\n"
            "ret\n"
            : : "r"(sp)
        );
    }
}

void task_info(void) {
    uart_puts("\nTask Status:\n");
    uart_puts("============\n");
    
    for(int i = 0; i < num_tasks; i++) {
        uart_puts("Task ");
        uart_puts(tasks[i].name);
        uart_puts(": ");
        
        switch(tasks[i].status) {
            case TASK_READY:
                uart_puts("READY");
                break;
            case TASK_RUNNING:
                uart_puts("RUNNING");
                break;
            case TASK_FINISHED:
                uart_puts("FINISHED");
                break;
        }
        uart_puts("\n");
    }
}
